"""知识检索服务"""

import os
import json
from typing import List, Dict, Any, Optional, Tuple
from pathlib import Path
import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
import numpy as np

from config.settings import config


class KnowledgeService:
    """知识检索服务"""
    
    def __init__(self, knowledge_dir: str = None, db_path: str = None):
        if knowledge_dir is None:
            knowledge_dir = Path(__file__).parent.parent.parent / "config" / "prompts" / "knowledge"
        if db_path is None:
            db_path = Path(__file__).parent.parent.parent / "data" / "chroma_db"
        
        self.knowledge_dir = Path(knowledge_dir)
        self.db_path = Path(db_path)
        self.db_path.mkdir(parents=True, exist_ok=True)
        
        # 初始化向量数据库
        self.client = chromadb.PersistentClient(
            path=str(self.db_path),
            settings=Settings(anonymized_telemetry=False)
        )
        
        # 初始化嵌入模型
        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        
        # 创建自定义嵌入函数
        def custom_embedding_function(texts):
            if isinstance(texts, str):
                texts = [texts]
            embeddings = self.embedding_model.encode(texts)
            return embeddings.tolist()
        
        # 获取或创建集合，使用自定义嵌入函数
        try:
            self.collection = self.client.get_collection(name="cairo_knowledge")
        except:
            # 如果集合不存在，创建新集合
            from chromadb.utils import embedding_functions
            
            class CustomEmbeddingFunction(embedding_functions.EmbeddingFunction):
                def __call__(self, input):
                    return custom_embedding_function(input)
            
            self.collection = self.client.create_collection(
                name="cairo_knowledge",
                embedding_function=CustomEmbeddingFunction(),
                metadata={"description": "Cairo and Starknet knowledge base"}
            )
        
        # 初始化知识库
        self._initialize_knowledge_base()

    def _initialize_knowledge_base(self):
        """初始化知识库"""
        # 检查是否已经有数据
        if self.collection.count() > 0:
            return
        
        # 加载并索引知识
        self._load_knowledge_files()
        self._load_builtin_knowledge()

    def _load_knowledge_files(self):
        """加载知识文件"""
        if not self.knowledge_dir.exists():
            return
        
        documents = []
        metadatas = []
        ids = []
        
        for file_path in self.knowledge_dir.glob("*.md"):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                
                # 分割文档为段落
                paragraphs = self._split_document(content)
                
                for i, paragraph in enumerate(paragraphs):
                    if paragraph.strip():
                        doc_id = f"{file_path.stem}_{i}"
                        documents.append(paragraph)
                        metadatas.append({
                            "source": str(file_path),
                            "type": "file",
                            "category": file_path.stem,
                            "paragraph_index": i
                        })
                        ids.append(doc_id)
            
            except Exception as e:
                print(f"Error loading {file_path}: {e}")
        
        if documents:
            self.collection.add(
                documents=documents,
                metadatas=metadatas,
                ids=ids
            )

    def _load_builtin_knowledge(self):
        """加载内置知识"""
        builtin_knowledge = [
            {
                "content": "Cairo 是一种为 STARK 证明系统设计的编程语言。它支持可证明计算，主要用于 Starknet 区块链上的智能合约开发。",
                "category": "basics",
                "tags": ["cairo", "starknet", "introduction"]
            },
            {
                "content": "felt252 是 Cairo 中的基本数据类型，表示一个 252 位的有限域元素。所有 Cairo 程序的基础都建立在 felt252 之上。",
                "category": "types",
                "tags": ["felt252", "types", "basics"]
            },
            {
                "content": "Cairo 中的函数使用 fn 关键字定义。函数可以有参数和返回值，支持泛型和 trait 约束。",
                "category": "functions",
                "tags": ["functions", "syntax", "fn"]
            },
            {
                "content": "Starknet 合约使用 #[contract] 属性标记。合约可以包含存储变量、事件定义和外部函数。",
                "category": "contracts",
                "tags": ["contracts", "starknet", "attributes"]
            },
            {
                "content": "Cairo 中的错误处理主要通过 Result<T, E> 类型和 panic! 宏实现。建议使用 Result 进行可恢复的错误处理。",
                "category": "error_handling",
                "tags": ["errors", "result", "panic"]
            },
            {
                "content": "Array<T> 是 Cairo 中的动态数组类型。可以使用 array! 宏创建数组，使用 .append() 方法添加元素。",
                "category": "collections",
                "tags": ["array", "collections", "data_structures"]
            },
            {
                "content": "trait 定义了类型必须实现的方法集合。impl 块用于为类型实现 trait 或添加方法。",
                "category": "traits",
                "tags": ["trait", "impl", "polymorphism"]
            },
            {
                "content": "Cairo 中的所有权系统确保内存安全。变量默认是不可变的，使用 mut 关键字声明可变变量。",
                "category": "ownership",
                "tags": ["ownership", "mutability", "safety"]
            },
            {
                "content": "Starknet 合约可以调用其他合约的函数。使用 IContractDispatcher 进行类型安全的合约调用。",
                "category": "contract_interaction",
                "tags": ["contract_calls", "dispatcher", "interaction"]
            },
            {
                "content": "Cairo 支持泛型编程，可以编写适用于多种类型的代码。泛型参数使用尖括号 <T> 声明。",
                "category": "generics",
                "tags": ["generics", "templates", "polymorphism"]
            }
        ]
        
        documents = []
        metadatas = []
        ids = []
        
        for i, item in enumerate(builtin_knowledge):
            documents.append(item["content"])
            metadatas.append({
                "source": "builtin",
                "type": "builtin",
                "category": item["category"],
                "tags": ",".join(item["tags"])
            })
            ids.append(f"builtin_{i}")
        
        self.collection.add(
            documents=documents,
            metadatas=metadatas,
            ids=ids
        )

    def _split_document(self, content: str) -> List[str]:
        """分割文档为段落，保持代码块和重要合约的完整性"""
        # 检查是否包含重要的合约代码（如ERC20）
        important_contracts = ['ERC20', 'MultiSigWallet', 'trait IERC20', 'mod ERC20']
        contains_important_contract = any(contract in content for contract in important_contracts)
        
        # 如果包含重要合约且不是太长，保持整个文档完整
        if contains_important_contract and len(content) < 5000:
            return [content]
        
        # 按双换行符分割
        paragraphs = content.split('\n\n')
        
        result = []
        i = 0
        while i < len(paragraphs):
            paragraph = paragraphs[i].strip()
            if not paragraph:
                i += 1
                continue
            
            # 检查是否是重要合约的开始
            is_contract_start = any(contract in paragraph for contract in important_contracts)
            
            if is_contract_start:
                # 尝试合并后续段落，直到找到完整的合约
                combined_content = paragraph
                j = i + 1
                
                # 继续合并段落，直到找到合约结束或达到合理长度
                while j < len(paragraphs) and len(combined_content) < 4000:
                    next_paragraph = paragraphs[j].strip()
                    if next_paragraph:
                        combined_content += '\n\n' + next_paragraph
                        
                        # 如果找到了合约结束标志，停止合并
                        if ('```' in next_paragraph and 
                            combined_content.count('```') % 2 == 0 and 
                            len(combined_content) > 1000):
                            break
                    j += 1
                
                result.append(combined_content)
                i = j + 1
            else:
                # 检查是否包含代码块
                contains_code_block = '```' in paragraph
                
                if contains_code_block:
                    # 对于代码块，保持完整性
                    if len(paragraph) > 2000:
                        # 尝试按代码块边界分割
                        code_sections = paragraph.split('```')
                        current_section = ""
                        for idx, section in enumerate(code_sections):
                            if idx % 2 == 0:  # 非代码部分
                                current_section += section
                            else:  # 代码部分
                                code_part = '```' + section + '```'
                                if len(current_section + code_part) > 2000 and current_section:
                                    result.append(current_section.strip())
                                    current_section = code_part
                                else:
                                    current_section += code_part
                        if current_section.strip():
                            result.append(current_section.strip())
                    else:
                        result.append(paragraph)
                else:
                    # 非代码段落，按原逻辑处理
                    if len(paragraph) > 500:
                        sentences = paragraph.split('。')
                        current_chunk = ""
                        for sentence in sentences:
                            if len(current_chunk + sentence) > 500 and current_chunk:
                                result.append(current_chunk.strip())
                                current_chunk = sentence
                            else:
                                current_chunk += sentence + "。"
                        if current_chunk.strip():
                            result.append(current_chunk.strip())
                    else:
                        result.append(paragraph)
                i += 1
        
        return result

    def search_knowledge(self, query: str, n_results: int = 5, 
                        category_filter: Optional[str] = None) -> List[Dict[str, Any]]:
        """搜索知识库"""
        try:
            # 构建查询条件
            where_clause = {}
            if category_filter:
                where_clause["category"] = category_filter
            
            # 执行搜索
            results = self.collection.query(
                query_texts=[query],
                n_results=n_results,
                where=where_clause if where_clause else None
            )
            
            # 格式化结果
            knowledge_items = []
            if results['documents'] and results['documents'][0]:
                for i, doc in enumerate(results['documents'][0]):
                    metadata = results['metadatas'][0][i] if results['metadatas'] else {}
                    distance = results['distances'][0][i] if results['distances'] else 1.0
                    
                    knowledge_items.append({
                        "content": doc,
                        "metadata": metadata,
                        "relevance_score": 1.0 - distance,  # 转换为相似度分数
                        "source": metadata.get("source", "unknown"),
                        "category": metadata.get("category", "general")
                    })
            
            return knowledge_items
        
        except Exception as e:
            print(f"Error searching knowledge: {e}")
            return []

    def get_relevant_knowledge(self, user_intent: str, code_context: Dict[str, Any], 
                             query: str) -> Tuple[List[Dict[str, Any]], List[float]]:
        """根据意图和上下文获取相关知识"""
        knowledge_items = []
        relevance_scores = []
        
        # 基于意图的知识检索
        # 特殊处理ERC20等常见合约类型
        if "erc20" in query.lower() or "代币合约" in query or "token contract" in query.lower():
            search_query = "ERC20 代币合约 token contract"
        else:
            intent_queries = {
                "qa": query,
                "code_analysis": f"代码分析 {code_context.get('code_type', '')} {query}",
                "error_diagnosis": f"错误诊断 错误处理 {query}",
                "deep_research": f"深入分析 原理 机制 {query}",
                "code_generation": f"代码示例 实现 {query}",
                "optimization": f"优化 性能 最佳实践 {query}",
                "tutorial": f"教程 指南 入门 {query}"
            }
            search_query = intent_queries.get(user_intent, query)
        
        # 执行搜索
        results = self.search_knowledge(search_query, n_results=8)
        
        # 如果有代码上下文，添加相关的代码知识
        if code_context.get("has_code"):
            code_type = code_context.get("code_type", "")
            if code_type:
                code_results = self.search_knowledge(f"{code_type} 语法 示例", n_results=3)
                results.extend(code_results)
        
        # 去重并排序
        seen_content = set()
        for item in results:
            content = item["content"]
            if content not in seen_content:
                seen_content.add(content)
                knowledge_items.append(item)
                relevance_scores.append(item["relevance_score"])
        
        # 按相关性排序
        sorted_items = sorted(zip(knowledge_items, relevance_scores), 
                            key=lambda x: x[1], reverse=True)
        
        if sorted_items:
            knowledge_items, relevance_scores = zip(*sorted_items)
            return list(knowledge_items)[:5], list(relevance_scores)[:5]
        
        return [], []

    def add_knowledge(self, content: str, category: str, metadata: Dict[str, Any] = None):
        """添加新知识"""
        if metadata is None:
            metadata = {}
        
        metadata.update({
            "category": category,
            "type": "user_added",
            "source": "manual"
        })
        
        doc_id = f"manual_{self.collection.count()}"
        
        self.collection.add(
            documents=[content],
            metadatas=[metadata],
            ids=[doc_id]
        )

    def update_knowledge_base(self):
        """更新知识库"""
        # 重新加载文件
        self._load_knowledge_files()


# 全局知识服务实例
knowledge_service = KnowledgeService()


def retrieve_knowledge(state: Dict[str, Any]) -> Dict[str, Any]:
    """增强的知识检索节点（集成外部工具）"""
    user_text = state["messages"][-1]["content"] if state["messages"] else ""
    user_intent = state.get("user_intent")
    cairo_context = state.get("cairo_context", {})
    
    # 1. 本地知识库检索
    local_knowledge, local_scores = knowledge_service.get_relevant_knowledge(
        user_intent.value if user_intent else "qa",
        cairo_context,
        user_text
    )
    
    # 2. 外部工具检索
    external_knowledge = []
    external_search_info = {"used_tools": [], "has_results": False}
    
    if config.external_tools.enable_external_search:
        try:
            from core.services.external_tools import external_tools_service
            import asyncio
            
            # 根据意图选择工具
            selected_tools = _select_tools_by_intent(user_intent, user_text, cairo_context)
            
            if selected_tools:
                print(f"Starting external search with tools: {selected_tools}")
                external_search_info["used_tools"] = selected_tools
                
                # 异步搜索外部工具
                try:
                    # 检查是否已经在事件循环中
                    try:
                        loop = asyncio.get_running_loop()
                        # 如果已经在事件循环中，使用线程池执行
                        import concurrent.futures
                        import threading
                        
                        def run_search():
                            new_loop = asyncio.new_event_loop()
                            asyncio.set_event_loop(new_loop)
                            try:
                                return new_loop.run_until_complete(
                                    external_tools_service.parallel_search(user_text, selected_tools)
                                )
                            finally:
                                new_loop.close()
                        
                        with concurrent.futures.ThreadPoolExecutor() as executor:
                            future = executor.submit(run_search)
                            external_results = future.result(timeout=30)
                            
                    except RuntimeError:
                        # 没有运行的事件循环，可以直接创建新的
                        loop = asyncio.new_event_loop()
                        asyncio.set_event_loop(loop)
                        try:
                            external_results = loop.run_until_complete(
                                external_tools_service.parallel_search(user_text, selected_tools)
                            )
                        finally:
                            loop.close()
                    
                    # 合并和排序结果
                    merged_results = external_tools_service.merge_and_rank_results(external_results, max_total=5)
                    
                    if merged_results:
                        external_search_info["has_results"] = True
                        print(f"External search found {len(merged_results)} results")
                    
                    # 转换为知识格式
                    for result in merged_results:
                        external_knowledge.append({
                            "content": result.content,
                            "metadata": {**result.metadata, "type": "external"} if result.metadata else {"type": "external"},
                            "source": result.source,
                            "category": "external",
                            "relevance_score": result.relevance_score,
                            "title": result.title,
                            "url": result.url
                        })
                        
                except Exception as search_error:
                    print(f"External search execution error: {search_error}")
                    external_knowledge = []
                    
        except Exception as e:
            print(f"External tools search error: {e}")
            external_knowledge = []
    
    print(f"Knowledge retrieval: local={len(local_knowledge)}, external={len(external_knowledge)}")
    
    # 3. 合并本地和外部知识
    all_knowledge = local_knowledge + external_knowledge
    all_scores = local_scores + [item["relevance_score"] for item in external_knowledge]
    
    # 4. 重新排序（本地知识优先，但高质量外部内容可以排前面）
    combined_items = list(zip(all_knowledge, all_scores))
    combined_items.sort(key=lambda x: (
        x[1] * (1.2 if x[0].get("metadata", {}).get("type") != "external" else 1.0)
    ), reverse=True)
    
    if combined_items:
        final_knowledge, final_scores = zip(*combined_items)
        final_knowledge = list(final_knowledge)[:8]  # 限制总数
        final_scores = list(final_scores)[:8]
    else:
        final_knowledge, final_scores = [], []
    
    # 更新状态
    state["retrieved_knowledge"] = final_knowledge
    state["knowledge_relevance_scores"] = final_scores
    state["external_search_info"] = external_search_info
    
    # 设置下一步动作
    if user_intent and hasattr(user_intent, 'value') and user_intent.value == "deep_research":
        state["next_action"] = "research_planning"
    else:
        state["next_action"] = "answer_generation"
    
    return state


def _select_tools_by_intent(user_intent, user_text: str, cairo_context: Dict[str, Any]) -> List[str]:
    """根据意图选择外部工具"""
    tools = []
    
    # 根据用户意图选择工具
    if user_intent:
        intent_value = user_intent.value if hasattr(user_intent, 'value') else str(user_intent)
        
        if intent_value == "error_diagnosis":
            # 错误诊断优先使用Stack Overflow和GitHub Issues
            tools.extend(["stackoverflow", "github"])
        elif intent_value == "code_analysis" or intent_value == "code_generation":
            # 代码相关优先使用GitHub和文档
            tools.extend(["github", "cairo_docs"])
        elif intent_value == "deep_research":
            # 深度研究使用所有工具
            tools.extend(["cairo_docs", "github", "stackoverflow"])
        else:
            # 默认使用文档搜索
            tools.append("cairo_docs")
    
    # 根据文本内容进一步筛选（这是关键的改进）
    text_lower = user_text.lower()
    
    # 如果提到错误、问题，添加Stack Overflow
    error_keywords = ["error", "错误", "问题", "bug", "issue", "fail", "panic", "debug", "报错", "出错"]
    if any(keyword in text_lower for keyword in error_keywords):
        if "stackoverflow" not in tools:
            tools.append("stackoverflow")
        print(f"Detected error keywords, added stackoverflow. Tools: {tools}")
    
    # 如果有代码块，添加GitHub搜索
    if cairo_context.get("has_code") or "```" in user_text:
        if "github" not in tools:
            tools.append("github")
        print(f"Detected code, added github. Tools: {tools}")
    
    # 如果没有选择任何工具，默认使用文档搜索
    if not tools:
        tools.append("cairo_docs")
    
    print(f"Final selected tools: {tools}")
    return tools[:3]  # 限制最多3个工具